#include <fstream>
#include <iostream>
#include <sstream>
#include <string>
using namespace std;

#include "Bitext.hh"
#include "Schedule.hh"
#include "Tokenizer.hh"

#include "Model.hh"
#include "UniformModel.hh"
#include "IBMModel1.hh"

ScheduleItem::ScheduleItem(string line){
  // tokenize line
  Tokenizer tokenizer(line);
  if(tokenizer.size() != 2){
    cerr << "incorrect schedule format .. each line: <model> <iterations>" << endl;
    exit(1);
  }

  // first model
  model = tokenizer[0];
  if(tokenizer[0] == "ibm_model1"){
  }else{
    cerr << model << " is not a supported model" << endl;
    exit(1);
  }

  // then iterations
  istringstream iss(tokenizer[1]);
  if(!(iss >> iterations)){
    cerr << tokenizer[1] << " is not a valid number of iterations" << endl;
    exit(1);
  }
}

Schedule::Schedule(string filename){
  // attempt to open file 
  ifstream file(filename.c_str());
  if(!file.is_open()){
    cerr << "could not open schedule file!" << endl;
    exit(1);
  }

  // while lines, add into schedule
  items.clear();
  for(string line; getline(file, line);){
    ScheduleItem item(line); 
    items.push_back(item);
  }
}

void Schedule::dump(){
  cout << "schedule dump!" << endl;
  for(list<ScheduleItem>::iterator i = items.begin(); i != items.end(); i++){
    cout << "model: " << i->model << ", iterations: " << i->iterations << endl;
  }
}

void Schedule::run(Bitext &bitext, Bitext &testset){
  // initialize to uniform model
  cout << "initializing to uniform model." << endl;
  Model *p_model = (Model*)new UniformModel();

  // run the schedule ..
  cout << "running schedule." << endl;
  for(list<ScheduleItem>::iterator i = items.begin(); i != items.end(); i++){
    //peek_expectation(bitext, p_model);
    bitext.inspect(p_model, 
		   bitext.get_filename() + ".trainset", 
		   bitext.get_filename() + ".trainset",
		   bitext.get_filename() + ".trainset",
		   -1);
    testset.inspect(p_model, 
		    bitext.get_filename() + ".testset",
		    bitext.get_filename() + ".testset", 
		    bitext.get_filename() + ".testset", 
		    -1);

    // request new model to transfer from old model.
    Model *p_new_model;
    if(i->model == "ibm_model1"){
      p_new_model = (Model*)new IBMModel1();
    }else{
      cerr << "internal error! (bad model name)" << endl;
      exit(1);
    }
    cout << "beginning transfer from " << p_model->get_model_name() << " to " << p_new_model->get_model_name() << "." << endl;
    p_new_model->transfer(bitext, p_model);

    // free old model.
    delete p_model;
    p_model = p_new_model;

    // do # iterations for new model.
    //peek_expectation(bitext, p_model);
    bitext.inspect(p_model, 
		   bitext.get_filename() + ".trainset", 
		   bitext.get_filename() + ".trainset",
		   bitext.get_filename() + ".trainset",
		   0);
    testset.inspect(p_model, 
		    bitext.get_filename() + ".testset",
		    bitext.get_filename() + ".testset",
		    bitext.get_filename() + ".testset",
		    0);
    cout << "running " << i->iterations << " iterations of model: " << i->model << endl;
    for(int j = 1; j <= i->iterations; j++){
      cout << "________________________iteration " << j << endl;
      p_model->iterate(bitext);

      //peek_expectation(bitext, p_model);
      bitext.inspect(p_model,
		     bitext.get_filename() + ".trainset",
		     bitext.get_filename() + ".trainset",
		     bitext.get_filename() + ".trainset",
		     j);
      testset.inspect(p_model,
		      bitext.get_filename() + ".testset",
		      bitext.get_filename() + ".testset",
		      bitext.get_filename() + ".testset",
		      j);
    }
  }
}

void Schedule::peek_expectation(Bitext &bitext, Model *p_model){
  // debug: let's have a look at the expected alignment of first sentence pair.
  cerr << "peek_expectation called!" << endl;
  SentencePair pair;
  bitext.rewind();
  while(bitext.next(pair)){
    cout << "(source) rows:    " << pair.source << endl;
    cout << "(target) columns: " << pair.target << endl;
    ExpectedCounts counts(pair);
    p_model->get_expected_counts(bitext, pair, counts);
    counts.print();
  }
  bitext.rewind();
}
